#!/usr/bin/env python

import roslib; roslib.load_manifest('team4')
import rospy
from pr2_tasks.tasks import Tasks
from pr2_tasks.pickplace import PickPlace
from pr2_tasks.tableware_detection import TablewareDetection
from pr2_tasks.pickplace_definitions import PlaceGoal
from pr2_python import arm_planner, transform_listener, gripper, world_interface, planning_scene_interface
from arm_mover import ArmMover
from object_manipulation_msgs.msg import Grasp, ClusterBoundingBox
from pr2_controllers_msgs.msg import JointTrajectoryAction
from geometry_msgs.msg import PointStamped, PoseStamped
from trajectory_msgs.msg import JointTrajectoryPoint, JointTrajectory
from tf.transformations import *
from sensor_msgs.msg import PointCloud

class Tray():
   def __init__(self):
      self.tray_location = PointStamped()
      self.last_grasps = None
      self.tray = None
      self.b_right = None
      self.b_left = None
      self.l_arm_mover = ArmMover('left_arm')
      self.r_arm_mover = ArmMover('right_arm')
      self.obj_bounding_box = None
      self.listener = transform_listener.get_transform_listener()
      self.drive_pose = PoseStamped()
      self.drive_pose.pose.position.x = 0.45
      self.drive_pose.pose.position.y = 0.0
      self.drive_pose.pose.position.z = 0.85
      self.drive_pose.pose.orientation.w = 1.0
      self.world = world_interface.WorldInterface()
      self.psi = planning_scene_interface.get_planning_scene_interface()


   def find_tray(self):
      det = TablewareDetection()
#      head_point = PointStamped()
#      head_point.header.frame_id = '/base_link'
#      head_point.point.x = 0.6
      res = det.detect_objects()
      
      for pg in res.pickup_goals:
         if pg.label == 'graspable': # limit this to unrecognized objects, in our context trays
            #self.obj_bounding_box = det._bounding_box_srv(pg.target.cluster) 
            self.tray = pg
      return

   def splitCluster(self):
      self.b_right = PointCloud()
      self.b_right.header = self.tray.target.cluster.header
      self.b_left = PointCloud()
      self.b_left.header = self.tray.target.cluster.header
      x_max = 0.0
      x_min = 100.0
      y_max = -100.0
      y_min = 100.0

      rospy.loginfo("Preparing to filter cluster of size {0}".format(len(self.tray.target.cluster.points)))
      # Find cluster stats
      for pt in self.tray.target.cluster.points:
         temp = PointStamped()
         temp.header = self.tray.target.cluster.header
         temp.point  = pt
         temp = self.listener.transformPoint("base_link", temp)
         if temp.point.x > x_max: 
            x_max = temp.point.x
         elif temp.point.x < x_min:
            x_min = temp.point.x
         if temp.point.y > y_max:
            y_max = temp.point.y
         elif temp.point.y < y_min:
            y_min = temp.point.y      
         pt = temp.point
          
      x_range = x_max - x_min
      x_center = (x_max+x_min)/2.0
      y_range = y_max - y_min
      y_center = (y_min+y_max)/2.0
      self.obj_bounding_box = [x_range, y_range]
      
      # Filter cluster
      for pt in self.tray.target.cluster.points:
         # if the ignored area in the middle is not wide enough, then make the y_range divide by a smaller number
         if pt.y >= y_center+(y_range/3.0) and ((pt.x <= x_center + (x_range/6.0)) and (pt.x >= x_center - (x_range/6.0))):
            self.b_left.points.append(pt)
         elif pt.y <= y_center-(y_range/3.0) and ((pt.x <= x_center + (x_range/6.0)) and (pt.x >= x_center - (x_range/6.0))): 
            self.b_right.points.append(pt)

      rospy.loginfo("Right cluster size: {0}".format(len(self.b_right.points)))
      rospy.loginfo("Left cluster size: {0}".format(len(self.b_left.points)))
      

   def gripClusters(self):
      pickplace = PickPlace()
      pg = self.tray
      pg.target.cluster = self.b_right
      pg.arm_name = 'right_arm'
      pg.lift.min_distance = 0.005
      pg.lift.desired_distance = 0.01
      pickplace.pickup(pg)
      # DOES NOT CATCH EXCEPTIONS

      pg.target.cluster = self.b_left
      pg.arm_name = 'left_arm'
      pg.lift.min_distance = 0.005
      pg.lift.desired_distance = 0.01
      pickplace.pickup(pg)
      # DOES NOT CATCH EXCEPTIONS
      
      # get coordiantes 
      
      # hack pose box
      print self.tray
      self.world.remove_collision_object(self.tray.collision_object_name)
      self.psi.reset()
      
      # add the object to the world in the gripper 
      pose = PoseStamped()
      pose.pose.orientation.w = 1.0
      r_gripper_pose = transform_listener.transform_pose("/base_link", '/r_wrist_roll_link', pose.pose )
      l_gripper_pose = transform_listener.transform_pose("/base_link", '/l_wrist_roll_link', pose.pose )
      y_gap = abs(r_gripper_pose.position.y - l_gripper_pose.position.y)
      
      right_pose = PoseStamped()
      right_pose.pose.orientation.w=1
      right_pose.header.frame_id='/base_link'
      right_pose.pose.position = r_gripper_pose.position
      right_pose.pose.position.y += y_gap*0.22
      right_pose.pose.position.z -= 0.15
      self.world.add_collision_box(right_pose, (0.4, y_gap*0.44, 0.05), 'tray_right_half')
      self.world.attach_object_to_gripper('right_arm','tray_right_half')
      self.psi.reset()

      left_pose = PoseStamped()
      left_pose.pose.orientation.w=1
      left_pose.header.frame_id='/base_link'
      left_pose.pose.position = l_gripper_pose.position
      left_pose.pose.position.y -= y_gap*0.22
      left_pose.pose.position.z -= 0.15
      self.world.add_collision_box(left_pose, (0.4, y_gap*0.44, 0.05), 'tray_left_half')
      self.world.attach_object_to_gripper('left_arm','tray_left_half')
      self.psi.reset()

   def getArmIK(self, pose=PoseStamped(), arm='right_arm'):
      arm_planner=[]
      if arm == 'right_arm':
	      arm_planner = self.r_arm_mover._arm_planner
      elif arm == 'left_arm':
	      arm_planner = self.l_arm_mover._arm_planner
      ik_sol = arm_planner.get_ik(pose)
      return ik_sol

   def moveTrayVertically(self, z_dist):
      pose = PoseStamped()
      pose.pose.orientation.w = 1
      target_pose = transform_listener.transform_pose("/base_link", '/l_wrist_roll_link', pose.pose)
      right_pose = transform_listener.transform_pose("/base_link", '/r_wrist_roll_link', pose.pose)
      target_pose.position.y = (target_pose.position.y+right_pose.position.y)/2.0
      target_pose.position.z += z_dist
      pose.pose = target_pose
      pose.header.frame_id = "/base_link"
      pose.header.stamp = rospy.Time(0)
      self.moveTray2Pose(pose)

   def moveTray2Pose2(self, p=PoseStamped()):
      # Calculate new pose relative to current pose
      l_ps = PoseStamped()
      l_ps.header.frame_id = "/base_link"
      l_ps.header.stamp = rospy.Time(0)
      l_ps.pose.orientation.w = 1.0
      r_ps = PoseStamped()
      r_ps.header.frame_id = "/base_link"
      r_ps.header.stamp = rospy.Time(0)
      r_ps.pose.orientation.w = 1.0
      new_l_ps = transform_listener.transform_pose("/base_link", '/l_wrist_roll_link', l_ps.pose)
      new_r_ps = transform_listener.transform_pose("/base_link", '/r_wrist_roll_link', r_ps.pose)
      arm_diff_y = abs(new_l_ps.position.y - new_r_ps.position.y)
   
      point = p.pose.position
      new_l_ps.position.x = point.x 
      new_l_ps.position.y = point.y + (arm_diff_y/2.0)
      new_l_ps.position.z = point.z
      new_r_ps.position.x = point.x
      new_r_ps.position.y = point.y - (arm_diff_y/2.0)
      new_r_ps.position.z = point.z
      l_ps.pose = new_l_ps
      r_ps.pose = new_r_ps

      # Create new plans
      r_arm_planner = self.r_arm_mover._arm_planner
      l_arm_planner = self.l_arm_mover._arm_planner
      r_traj = r_arm_planner.plan_interpolated_ik(r_ps, collision_aware=False)
      # DOES NOT CATCH EXCEPTIONS
      l_traj = l_arm_planner.plan_interpolated_ik(l_ps, collision_aware=False)
      # DOES NOT CATCH EXCEPTIONS

      # Splice plans into one trajectory
      spliced = JointTrajectory()
      spliced.joint_names = r_traj.joint_names + l_traj.joint_names
      m = min(len(r_traj.points), len(l_traj.points))
      for i in range(0,m):
         rospy.loginfo("Splicing point {0}".format(i))
         spliced.points.append(JointTrajectoryPoint())
         spliced.points[i].positions = r_traj.points[i].positions + l_traj.points[i].positions
         spliced.points[i].velocities = r_traj.points[i].velocities + l_traj.points[i].velocities
         spliced.points[i].time_from_start = r_traj.points[i].time_from_start

      # Exec two arm trajectory
      rospy.loginfo("Executing spliced trajectory")
      self.l_arm_mover.execute_two_arm_trajectory(spliced)
      # DOES NOT CATCH EXCEPTIONS
      rospy.loginfo("Executed spliced trajectory")
      
   def moveTray2Pose(self, p=PoseStamped()):
      # Calculate new pose relative to current pose
      self.psi.reset()
      l_ps = PoseStamped()
      l_ps.header.frame_id = "/base_link"
      l_ps.header.stamp = rospy.Time(0)
      l_ps.pose.orientation.w = 1.0
      r_ps = PoseStamped()
      r_ps.header.frame_id = "/base_link"
      r_ps.header.stamp = rospy.Time(0)
      r_ps.pose.orientation.w = 1.0
      new_l_ps = transform_listener.transform_pose("/base_link", '/l_wrist_roll_link', l_ps.pose)
      new_r_ps = transform_listener.transform_pose("/base_link", '/r_wrist_roll_link', r_ps.pose)
      arm_diff_x = abs(new_l_ps.position.x - new_r_ps.position.x)
      arm_diff_y = abs(new_l_ps.position.y - new_r_ps.position.y)
      arm_diff_z = abs(new_l_ps.position.z - new_r_ps.position.z)
      
      trayPose = PoseStamped()
      trayPose.pose.position.x = new_r_ps.position.x + 0.5*arm_diff_x
      trayPose.pose.position.y = new_r_ps.position.y + 0.5*arm_diff_y
      trayPose.pose.position.z = new_r_ps.position.z + 0.5*arm_diff_z
      trayPose.pose.orientation.w = 1.0
      trayPose.header.frame_id = "/base_link"
      
      tray_diff_x = p.pose.position.x - trayPose.pose.position.x
      tray_diff_y = p.pose.position.y - trayPose.pose.position.y
      tray_diff_z = p.pose.position.z - trayPose.pose.position.z
   
      point = p.pose.position
      new_l_ps.position.x += tray_diff_x
      new_l_ps.position.y += tray_diff_y
      new_l_ps.position.z += tray_diff_z
      new_r_ps.position.x += tray_diff_x
      new_r_ps.position.y += tray_diff_y
      new_r_ps.position.z += tray_diff_z
      l_ps.pose = new_l_ps
      r_ps.pose = new_r_ps

      # Create new plans
      r_arm_planner = self.r_arm_mover._arm_planner
      l_arm_planner = self.l_arm_mover._arm_planner
      r_traj = r_arm_planner.plan_interpolated_ik(r_ps, collision_aware=False)
      # DOES NOT CATCH EXCEPTIONS
      l_traj = l_arm_planner.plan_interpolated_ik(l_ps, collision_aware=False)
      # DOES NOT CATCH EXCEPTIONS

      # Splice plans into one trajectory
      spliced = JointTrajectory()
      spliced.joint_names = r_traj.joint_names + l_traj.joint_names
      m = min(len(r_traj.points), len(l_traj.points))
      for i in range(0,m):
         rospy.loginfo("Splicing point {0}".format(i))
         spliced.points.append(JointTrajectoryPoint())
         spliced.points[i].positions = r_traj.points[i].positions + l_traj.points[i].positions
         spliced.points[i].velocities = r_traj.points[i].velocities + l_traj.points[i].velocities
         spliced.points[i].time_from_start = r_traj.points[i].time_from_start

      # Exec two arm trajectory
      rospy.loginfo("Executing spliced trajectory")
      self.l_arm_mover.execute_two_arm_trajectory(spliced)
      # DOES NOT CATCH EXCEPTIONS
      rospy.loginfo("Executed spliced trajectory")

   def fake_pick_up_tray(self):
      pose = PoseStamped()
      pose.pose.orientation.w = 1.0
      r_gripper_pose = transform_listener.transform_pose("/base_link", '/r_wrist_roll_link', pose.pose )
      l_gripper_pose = transform_listener.transform_pose("/base_link", '/l_wrist_roll_link', pose.pose )
      y_gap = abs(r_gripper_pose.position.y - l_gripper_pose.position.y)

      right_pose = PoseStamped()
      right_pose.pose.orientation.w=1
      right_pose.header.frame_id='/base_link'
      right_pose.pose.position = r_gripper_pose.position
      right_pose.pose.position.y += y_gap*0.22
      right_pose.pose.position.z -= 0.15
      self.world.add_collision_box(right_pose, (0.4, y_gap*0.44, 0.05), 'tray_right_half')
      self.world.attach_object_to_gripper('right_arm','tray_right_half')

      left_pose = PoseStamped()
      left_pose.pose.orientation.w=1
      left_pose.header.frame_id='/base_link'
      left_pose.pose.position = l_gripper_pose.position
      left_pose.pose.position.y -= y_gap*0.22
      left_pose.pose.position.z -= 0.15
      self.world.add_collision_box(left_pose, (0.4, y_gap*0.44, 0.05), 'tray_left_half')
      self.world.attach_object_to_gripper('left_arm','tray_left_half')

      self.psi.reset()
      self.moveTrayVertically(0.1)
      #self.moveTray2Pose(self.drive_pose)
      
   def pick_up_tray(self):
      self.find_tray()
      self.splitCluster()
      self.gripClusters()
      self.moveTrayVertically(0.2)
      #self.moveTray2Pose(self.drive_pose)
      
   def put_down_tray(self, pose=PoseStamped()):
      pose.pose.position.z += 0.1
      self.moveTray2Pose(pose)
      self.moveTrayVertically(-0.1)
      # detatch 
      self.world.detach_object('right_arm', 'tray_right_half')
      self.world.remove_collision_object('tray_right_half')
      self.world.detach_object('left_arm', 'tray_left_half')
      self.world.remove_collision_object('tray_left_half')
      # open gripper
      gripper.Gripper('right_arm').open()
      gripper.Gripper('left_arm').open()
      self.moveTrayVertically(0.1)

#   def put_down_tray(self):
#      self.moveTrayVertically(-0.1)
#      # Store location
#      rospy.loginfo("Waiting on transform")
#      transform_listener.get_transform_as_tr('/base_link', '/r_wrist_roll_link')
#      rospy.loginfo("Found transform")
#
#      # Calculate new pose relative to current pose
#      l_gripper = PoseStamped()
#      r_gripper = PoseStamped()
#      l_gripper.header.frame_id = "/base_link"
#      l_gripper.header.stamp = rospy.Time(0)
#      l_gripper.pose.orientation.w = 1.0
#      l_gripper.pose = transform_listener.transform_pose("/base_link", '/l_wrist_roll_link', l_gripper.pose)
#      r_gripper.pose = transform_listener.transform_pose("/base_link", '/r_wrist_roll_link', l_gripper.pose)
#      l_gripper.pose.position.x = (l_gripper.pose.position.x + r_gripper.pose.position.x)/2.0
#      l_gripper.pose.position.y = (l_gripper.pose.position.y + r_gripper.pose.position.y)/2.0
#      l_gripper.pose.position.z = (l_gripper.pose.position.z + r_gripper.pose.position.z)/2.0
#      self.tray_location = l_gripper
# 
#   def get_tray_out_of_the_way(self):
#      self.moveTrayVertically(-0.5)
#   def move_tray_back_up(self):
#      self.moveTrayVertically(0.5)

 #
#if __name__ == "__main__":
   #rospy.init_node("move_both_arms_test")
   #test = Tray()
     #
   #test.pick_up_tray()
   #
   #rospy.sleep(2.0)
   #test.put_down_tray()
     
